
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include "../iface/debug.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "DSAPublicKey.h"
#include "DSAPrivateKey.h"

/** Constructor for a blank DSA private key. */
DSAPrivateKey::DSAPrivateKey()
{
   key = NULL;
   keyname = NULL;
}

/** Constructor for a blank DSA private key with a name.
 * 
 * @param name
 * The name to attach to the blank key. */
DSAPrivateKey::DSAPrivateKey(const char *name)
{
   key = NULL;
   keyname = NULL;
   setKeyName(name);
}

/** Copy constructor for an DSA private key.
 * 
 * @param other
 * The DSA private key to copy from. */
DSAPrivateKey::DSAPrivateKey(const DSAPrivateKey &other)
{
   key = DSA_new();
   if (key == NULL)
   {
      errprint("DSA_new failed");
      return;
   }
   copyDSASSLPrivateKey(key, other.key);
   keyname = NULL;
   setKeyName(other.keyname);
}

/** Construct the DSA private key from a PEM encoded PKCS #8 text format
 * DSA private key.
 * 
 * @param pemtext
 * The PEM encoded PKCS #8 text that represents an DSA private key, from
 * which this key is to be constructed.
 * 
 * @param name
 * The name to be attached to the key, once constructed. */
DSAPrivateKey::DSAPrivateKey(const char *pemtext, const char *name)
{
   importKeyFromASCII(pemtext, name);
}

/** Construct the DSA private key from an encrypted and passworded PEM
 * encoded PKCS #8 text format DSA private key.
 * 
 * @param pemtext
 * The PEM encoded PKCS #8 text that represents an DSA public key, from
 * which this key is to be constructed.
 * 
 * @param password
 * The password to decrypt the PEM encoded PKCS #8 text that
 * represents the DSA public key.
 * 
 * @param name
 * The name to be attached to the key, once constructed. */
DSAPrivateKey::DSAPrivateKey(const char *pemtext, const char *password, const char *name)
{
   importKeyFromASCIIWithPassword(pemtext, password, name);
}

/** Construct the DSA private key from an DSA private key in SSL format.
 * 
 * @param prvkey
 * The SSL format DSA private key to construct this key from.
 * 
 * @param name
 * The name to be attached to the key, once constructed. */
DSAPrivateKey::DSAPrivateKey(DSA *prvkey, const char *name)
{
   key = DSA_new();
   if (key == NULL)
   {
      errprint("DSA_new failed");
      return;
   }
   copyDSASSLPrivateKey(key, prvkey);
   keyname = NULL;
   setKeyName(name);
}

/** Copy assignment operator for an DSA private key.
 * 
 * @param other
 * The DSA private key to copy from. */
DSAPrivateKey &DSAPrivateKey::operator=(const DSAPrivateKey &other)
{
   key = DSA_new();
   if (key == NULL)
   {
      errprint("DSA_new failed");
      return *this;
   }
   copyDSASSLPrivateKey(key, other.key);
   keyname = NULL;
   setKeyName(other.keyname);
   return *this;
}

/** Destructor for the DSA private key. */
DSAPrivateKey::~DSAPrivateKey()
{
   if (key) DSA_free(key);
   key = NULL;
   if (keyname) free((void *)keyname);
   keyname = NULL;
}

/** Obtains a cryptlib context which has had the key loaded into it.
 * 
 * @param[out] keycontext
 * A pointer to a cryptlib CRYPT_CONTEXT which is to be initialised with the
 * key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int DSAPrivateKey::getCryptlibContext(CRYPT_CONTEXT *keycontext)
{
   int status;
   CRYPT_PKCINFO_DLP *dsakey = (CRYPT_PKCINFO_DLP *)malloc(sizeof(CRYPT_PKCINFO_DLP));
   if (getCryptlibKeyComponents(dsakey) != 0)
   {
      free(dsakey);
      return -1;
   }
   int magictriescntxt = 0;
   do
   {
      status = cryptCreateContext(keycontext, CRYPT_UNUSED, CRYPT_ALGO_DSA);
      if (status == CRYPT_ERROR_TIMEOUT) apr_sleep(rand() % CRYPT_MAX_TIMESLEEP);
      magictriescntxt++;
   } while ((status == CRYPT_ERROR_TIMEOUT) && (magictriescntxt < CRYPT_MAX_RETRIES));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keycontext, CRYPT_CTXINFO_LABEL, keyname, strlen(keyname));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keycontext, CRYPT_CTXINFO_KEY_COMPONENTS, dsakey, sizeof(CRYPT_PKCINFO_DLP));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   cryptDestroyComponents(dsakey);
   free(dsakey);
   return 0;
failover:
   cryptDestroyComponents(dsakey);
   free(dsakey);
   return -1;
}

/** Obtains a cryptlib context meant particularly for signing.
 * 
 * @param[out] keycontext
 * A pointer to a cryptlib CRYPT_CONTEXT which is to be initialised with the
 * key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int DSAPrivateKey::getCryptlibSigningContext(CRYPT_CONTEXT *keycontext)
{
   return getCryptlibContext(keycontext);
}

/** Obtains a set of key components that can be loaded into cryptlib as a
 * private key.
 * 
 * @param[out] dsakey
 * A pointer to a cryptlib DSA key components structure which is to be
 * initialised with the key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int DSAPrivateKey::getCryptlibKeyComponents(CRYPT_PKCINFO_DLP *dsakey)
{
   unsigned char *buffera[5];
   cryptInitComponents(dsakey, CRYPT_KEYTYPE_PRIVATE);
   buffera[0] = (unsigned char *)malloc(BN_num_bytes(key->p));
   if (!(buffera[0])) { errprint("Could not allocate buffer for key->p."); return -1; }
   buffera[1] = (unsigned char *)malloc(BN_num_bytes(key->q));
   if (!(buffera[1])) { errprint("Could not allocate buffer for key->q."); return -1; }
   buffera[2] = (unsigned char *)malloc(BN_num_bytes(key->g));
   if (!(buffera[2])) { errprint("Could not allocate buffer for key->g."); return -1; }
   buffera[3] = (unsigned char *)malloc(BN_num_bytes(key->priv_key));
   if (!(buffera[3])) { errprint("Could not allocate buffer for key->priv_key (x)."); return -1; }
   buffera[4] = (unsigned char *)malloc(BN_num_bytes(key->pub_key));
   if (!(buffera[4])) { errprint("Could not allocate buffer for key->pub_key (y)."); return -1; }
   BN_bn2bin(key->p, buffera[0]);
   BN_bn2bin(key->q, buffera[1]);
   BN_bn2bin(key->g, buffera[2]);
   BN_bn2bin(key->priv_key, buffera[3]);
   BN_bn2bin(key->pub_key, buffera[4]);
   cryptSetComponent(dsakey->p, buffera[0], BN_num_bits(key->p));
   cryptSetComponent(dsakey->q, buffera[1], BN_num_bits(key->q));
   cryptSetComponent(dsakey->g, buffera[2], BN_num_bits(key->g));
   cryptSetComponent(dsakey->x, buffera[3], BN_num_bits(key->priv_key));
   cryptSetComponent(dsakey->y, buffera[4], BN_num_bits(key->pub_key));
   for (int i = 0; i < 5; i++) free(buffera[i]);
   return 0;
}

/** Get a copy of the private key as an SSL format object.
 * 
 * @param[out] prvkey
 * A pointer to a place to put a pointer to the new SSL format DSA private key
 * which is a copy of this key. The responsibility for freeing the new
 * DSA in SSL format is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::getSSLKeyComponents(DSA **prvkey)
{
   *prvkey = DSA_new();
   if ((*prvkey) == NULL)
   {
      errprint("DSA_new failed");
      return -1;
   }
   return copyDSASSLPrivateKey(*prvkey, key);
}

/** Get a copy of the private key in PEM encoded PKCS #8 text format.
 * 
 * @param[out] pemtext
 * A pointer to a place to put a pointer to the exported key in PEM
 * encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::getPEMKeyData(char **pemtext)
{
   char *data;
   int damt;
   int retval = 0;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_DSA(epkey, key);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PrivateKey(bo, epkey, 0, 0, 0, 0, 0))
   {
      retval = 1;
      errprint("could not export private key");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   *pemtext = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(*pemtext, data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
//   printf("%s\n", dbuf);
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

/** Get a copy of the public key in PEM encoded and encrypted PKCS #8 text
 * format, protected with the given password.
 * 
 * @param[out] pemtext
 * A pointer to a place to put a pointer to the exported key in encrypted
 * PEM encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @param password
 * The password to encrypt and decrypt the key with.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::getPEMKeyData(char **pemtext, const char *password)
{
   char *data;
   int damt;
   int retval = 0;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_DSA(epkey, key);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PKCS8PrivateKey(bo, epkey, EVP_des_ede3_cbc(), NULL, 0, 0, (void *)password))
   {
      retval = 1;
      errprint("could not export private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   (*pemtext) = (char *)malloc(sizeof(char)*(damt+1));
   memcpy((*pemtext), data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

/** Compare this PrivateKey to a PublicKey, and see if they are part of the
 * same key pair (one decrypts any content encrypted by the other).
 * 
 * @param pubkey
 * The public key to compare to this private key to see if they make a valid
 * pair.
 * 
 * @return
 * 1 if they make a matching pair of keys, 0 if they do not, and negative
 * values represent errors. */
int DSAPrivateKey::compareKeyBelongs(PublicKey *pubkey)
{
   if (pubkey == NULL)
   {
      errprint("comparing against NULL key");
      return 0;
   }
   switch (pubkey->getKeyType())
   {
      case KEY_TYPE_PUBLIC_DSA:
      {
         DSA *dsaKeyOther = NULL;
         ((DSAPublicKey *)pubkey)->getSSLKeyComponents(&dsaKeyOther);
         EVP_PKEY *evp1 = EVP_PKEY_new();
         EVP_PKEY *evp2 = EVP_PKEY_new();
         EVP_PKEY_set1_DSA(evp1, key);
         EVP_PKEY_set1_DSA(evp2, dsaKeyOther);
         int retval = EVP_PKEY_cmp(evp1, evp2);
         EVP_PKEY_free(evp1);
         EVP_PKEY_free(evp2);
         DSA_free(dsaKeyOther);
         fSSLCheckError();
         switch (retval)
         {
            case -2:
               errprint("key compare operation not supported (EVP_PKEY_cmp())");
               break;
            case -1:
               errprint("tried to compare different key types (EVP_PKEY_cmp())");
               return 0;
            case 0:
               return 0;
            case 1:
               return 1;
         }
         return 0;
         break;
      }
      default:
      {
         return 0;
         break;
      }
   }
}

/** Copy an SSL format private key.
 * 
 * @param[out] dst
 * Pointer to an initialised blank SSL format DSA key to copy into.
 * 
 * @param src
 * Pointer to an existing SSL format DSA private key to copy from.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::copyDSASSLPrivateKey(DSA *dst, DSA *src)
{
   if (dst == NULL) { errprint("NULL private key destination"); return -1; }
   if (src == NULL) { errprint("NULL private key source"); return 1; }
   // Copy public key
   dst->p = BN_new();
   if (BN_copy(dst->p, src->p) != (dst->p)) { errprint("Could not copy prv->p"); return -1; }
   dst->q = BN_new();
   if (BN_copy(dst->q, src->q) != (dst->q)) { errprint("Could not copy prv->q"); return -1; }
   dst->g = BN_new();
   if (BN_copy(dst->g, src->g) != (dst->g)) { errprint("Could not copy prv->g"); return -1; }

   // Copy private key
   dst->priv_key = BN_new();
   if (BN_copy(dst->priv_key, src->priv_key) != (dst->priv_key)) { errprint("Could not copy prv->priv_key (x)"); return -1; }
   dst->pub_key = BN_new();
   if (BN_copy(dst->pub_key, src->pub_key) != (dst->pub_key)) { errprint("Could not copy prv->pub_key (y)"); return -1; }
   return 0;
}

/** Import and name an DSA private key from an DSA private key in encrypted
 * PEM format PKCS #8.
 * 
 * @param pemtext
 * The PEM encoded PKCS #8 text that represents an DSA public key, from
 * which this key is to be constructed.
 * 
 * @param password
 * The password to decrypt the PEM encoded PKCS #8 text that
 * represents the DSA public key.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::importPrivateKeyWithOptionalPassword(const char *pemtext, const char *password)
{
   int retval = 0;
   BIO *bo = BIO_new_mem_buf((void *)pemtext, -1);
   EVP_PKEY *epkey = PEM_read_bio_PrivateKey(bo, NULL, 0, (void *)password);
   if (!epkey)
   {
      retval = 1;
      errprint("could not import private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   if (key) DSA_free(key);
   key = NULL;
   key = EVP_PKEY_get1_DSA(epkey);
   if (key == NULL)
   {
      retval = -1;
      errprint("failed private key import");
   }
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

/** Import an DSA private key from some PEM format PKCS #8 text.
 * 
 * @param text
 * The PEM encoded PKCS #8 text that represents an DSA public key, from
 * which this key is to be constructed.
 * 
 * @param name
 * The name to be attached to the key.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::importKeyFromASCII(const char *text, const char *name)
{
   return importKeyFromASCIIWithPassword(text, NULL, name);
}

/** Import and name an DSA private key from an DSA private key in encrypted
 * PEM format PKCS #8.
 * 
 * @param text
 * The PEM encoded PKCS #8 text that represents an DSA public key, from
 * which this key is to be constructed.
 * 
 * @param password
 * The password to decrypt the PEM encoded PKCS #8 text that
 * represents the DSA public key.
 * 
 * @param name
 * The name to be attached to the key.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::importKeyFromASCIIWithPassword(const char *text, const char *password, const char *name)
{
   setKeyName(name);
   int retval = importPrivateKeyWithOptionalPassword(text, NULL);
   return retval;
}

/** Export an DSA private key to some PEM format PKCS #8 text.
 * 
 * @param[out] text
 * A pointer to a place to put a pointer to the exported key in PEM
 * encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::exportKeyToASCII(char **text)
{
   return getPEMKeyData(text);
}

/** Export an DSA private key to some PEM format PKCS #8 text encrypted with
 * the given password.
 * 
 * @param[out] text
 * A pointer to a place to put a pointer to the exported key in encrypted
 * PEM encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @param password
 * The password to encrypt and decrypt the key with.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int DSAPrivateKey::exportKeyToASCIIWithPassword(char **text, const char *password)
{
   return getPEMKeyData(text, password);
}
